#!/usr/bin/env python

# *************************************************************************

# Copyright 2010 Texas A&M University, 
#                College Station, Texas
#                Contact: R Crosby (rwc@cs.tamu.edu)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# *************************************************************************

import sys
from QQPkg import QSub

# *************************************************************************

def ScriptProlog(sc, pd, config, opt):
    """Common job prolog information"""

    if config.has_section('mail'):
        sc.WriteCmd('-m', config.get('mail', 'when'))
        sc.WriteCmd('-M', config.get('mail', 'to'))
    
    batch = opt['batch']
    
    sc.WriteCmd('-l', 'walltime={0}'.format(config.get(batch, 'walltime')))
    sc.WriteCmd('-l', 'nodes={0}'.format(config.get(batch, 'nodes')))

    sc.WriteCmd('-o', '{0}/{1}.${{PBS_JOBID}}.stdout'.format(pd, sc.jobname))
    sc.WriteCmd('-e', '{0}/{1}.${{PBS_JOBID}}.stderr'.format(pd, sc.jobname))

    sc.WriteLn('') 
    sc.WriteLn('#!/bin/bash')
    sc.WriteLn('') 
    sc.WriteLn('WORK={0}'.format(pd))
    sc.WriteLn('') 

# *************************************************************************

def SubmitProcessingJobs ( nt, segf, pd, config, opt ):

    njobs = (len(segf)**2 + len(segf)) /2             # (n + 1) choose 2
    
    sc = QSub.QSubScript(pd + '/ExecScript.pbs', 'pqq', opt['queue'])

    sc.WriteCmd('-t', '0-{0}'.format(njobs))

    ScriptProlog(sc, pd, config, opt)

    i = 0
    for y in range(0, len(segf)):
        for x in range(y, len(segf)):

            xs = 0
            tr = min(opt['segsize'], nt-(y*opt['segsize']))
            ys = 0 if x == y else tr

            sc.WriteTx('OPT[{0}]="--xs={1} --ys={2} --trees={3} --bin-file=${{WORK}}/{4}_{0:0>5d}.bin'.
                       format(i, xs, ys, tr, sc.jobname))

            sc.WriteTx("${{WORK}}/{0}".format(segf[x]))
            
            if x <> y:
                sc.WriteTx("${{WORK}}/{0}".format(segf[y]))

            sc.WriteLn('"') 

            i += 1
    
#    OPT[0]='--xs=0 --ys=0 --trees=10 tf1'
#    OPT[1]='--xs=0 --ys=10 --trees=10 tf1 tf2'
                
    sc.WriteLn('') 
    sc.WriteLn('pushd {0}'.format(config.get('directories', 'output')))
    sc.WriteLn('{0}/qq {1}${{OPT[${{PBS_ARRAYID}}]}}'.
               format(sys.path[0],
                      '-v ' if opt['verbose'] else ''))
    sc.WriteLn('popd')

    jobno = sc.Submit()

    print 'Processing Job(s) {0} submitted from script {1}'.format(jobno, sc.filename)

    return sc.jobname, jobno, njobs

# *************************************************************************

def SubmitBackendJob ( jnm, jno, nj, pd, nsf, config, opt ):

    sc = QSub.QSubScript(pd + '/BackendScript.pbs', 'pqqback', opt['queue'])

    sc.WriteCmd('-W', 'depend=afteranyarray:%(jn)s' % {'jn' : jno})

    ScriptProlog(sc, pd, config, opt)

    sc.WriteLn('pushd '+config.get('directories', 'output'))
    sc.WriteLn('{0}/pqqback {1}--work-dir=${{WORK}} --job-count={2} --job-name={3} --job-id={4} --segments={5}'.
               format(sys.path[0],
                      '-v ' if opt['verbose'] else '', 
                      nj, 
                      jnm,
                      jno,
                      nsf))
    sc.WriteLn('popd')

    jobno = sc.Submit()

    print 'Backend Job {0} submitted from script {1}'.format(jobno, sc.filename)

    return sc.jobname, jobno

# *************************************************************************

if __name__ == "__main__":
    print 'Needs to be run from QQFront main'
